package benchmarks.tpch.mongo;

import api.connection.DBConnection;
import api.result.Result;
import api.result.ResultKV;
import api.result.UnitResult;
import benchmarks.tpch.TPCHLoader;
import benchmarks.tpch.mongo.domain.*;
import benchmarks.tpch.mongo.dto.*;
import benchmarks.tpch.utils.DBtool;
import benchmarks.tpch.utils.Tuple;
import com.mongodb.*;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.model.Indexes;
import org.bson.Document;
import org.mongodb.morphia.Morphia;

import java.io.*;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import static benchmarks.tpch.utils.Types.*;

public class TPCHLoaderMongoFast extends TPCHLoader{
    private DBConnection<MongoClient> dbConnection;
    private MongoClient mongoClient;
    private final static int configCommitCount = 10000; // commit every n records
    private Morphia morphia;

    public TPCHLoaderMongoFast(double scaleFactor) {
        super(scaleFactor);
    }

    @Override
    public Result loadData(DBConnection connection) {
        long start = System.currentTimeMillis();

        DBtool.prepareTBL(getScaleFactor());
        System.out.println("start load tables");
        loadTBL();
//        System.out.println("start refactor schema");
//        refactorSchema();
        System.out.println("start create database");
        this.dbConnection = connection;
        this.mongoClient = dbConnection.makeConnection();
        morphia = new Morphia();
        morphia.mapPackage("benchmarks.tpch.mongo.dto");
        mongoClient.getDatabase("tpch").getCollection("deals").drop();
        loadLineItemTable();
//        int recordNum = insertToMongo();
        System.out.println("start build index");
        createIndex();
        long end = System.currentTimeMillis();
        ResultKV<UnitResult> resTime = new ResultKV<>(
                "tpc-h\tMongo\t" + getScaleFactor() + "\tloadData",
                new UnitResult(end - start));
        ResultKV<UnitResult> resRecord = new ResultKV<>(
                "tpc-h\tMongo\t" + getScaleFactor() + "\ttotalRecords",
                new UnitResult(totalRows, "records"));
        CommandResult status = mongoClient.getDB("tpch").getCollection("deals").getStats();
        ResultKV<UnitResult> resSize = new ResultKV<>(
                "tpc-h\tMongo\t" + getScaleFactor() + "\tdataSize",
                new UnitResult(status.getInt("size") >> 10, "kb"));
        ResultKV<UnitResult> indexSize = new ResultKV<>(
                "tpc-h\tMongo\t" + getScaleFactor() + "\tindexSize",
                new UnitResult(status.getInt("totalIndexSize") >> 10, "kb"));
        System.out.println("loading tpc-h on mongo ok");
        mongoClient.close();
        customerHashMap.clear();
        nationHashMap.clear();
        ordersHashMap.clear();
        partHashMap.clear();
        partSuppHashMap.clear();
        regionHashMap.clear();
        supplierHashMap.clear();
        return resTime.union(resRecord).union(resSize).union(indexSize);
    }

    private void createIndex() {
        MongoCollection<Document> collection = mongoClient.getDatabase("tpch").getCollection("deals");
        collection.createIndex(Indexes.hashed("order.orderkey"));
        collection.createIndex(Indexes.hashed("partsupp.part.partkey"));
        collection.createIndex(Indexes.hashed("partsupp.supplier.suppkey"));
        collection.createIndex(Indexes.ascending("order.orderdate"));
        collection.createIndex(Indexes.ascending("shipdate"));
        collection.createIndex(Indexes.ascending("commitdate"));
        collection.createIndex(Indexes.ascending("receiptdate"));
    }

//    private int insertToMongo() {
//        final Morphia morphia = new Morphia();
//        morphia.mapPackage("benchmarks.tpch.mongo.dto");
//        DB db = mongoClient.getDB("tpch");
//        DBCollection collection = db.getCollection("deals");
//        collection.drop();
//        int count = 0;
//        List<DBObject> todoList = new ArrayList<>();
//        for (LineItemDto lineItemDto : data) {
//            todoList.add(morphia.toDBObject(lineItemDto));
//            if (++count % configCommitCount == 0) {
//                collection.insert(todoList);
//                todoList.clear();
//            }
//        }
//        if (todoList.size() > 0) {
//            collection.insert(todoList);
//        }
//        return count;
//    }

    private LineItemDto transLineItem(LineItem lineItem) {
        LineItemDto lineItemDto = new LineItemDto();
        lineItemDto.setLinenumber(lineItem.getL_linenumber());
        lineItemDto.setQuantity(lineItem.getL_quantity());
        lineItemDto.setExtendedprice(lineItem.getL_extendedprice());
        lineItemDto.setDiscount(lineItem.getL_discount());
        lineItemDto.setTax(lineItem.getL_tax());
        lineItemDto.setReturnflag(lineItem.getL_returnflag());
        lineItemDto.setLinestatus(lineItem.getL_linestatus());
        lineItemDto.setShipdate(lineItem.getL_shipdate());
        lineItemDto.setCommitdate(lineItem.getL_commitdate());
        lineItemDto.setReceiptdate(lineItem.getL_receiptdate());
        lineItemDto.setShipinstruct(lineItem.getL_shipinstruct());
        lineItemDto.setShipmode(lineItem.getL_shipmode());
        lineItemDto.setComment(lineItem.getL_comment());

        OrderDto orderDto = new OrderDto();
        Orders order = ordersHashMap.get(lineItem.getL_orderkey());
        orderDto.setOrderkey(order.getO_orderkey());
        orderDto.setOrderstatus(order.getO_orderstatus());
        orderDto.setTotalprice(order.getO_totalprice());
        orderDto.setOrderdate(order.getO_orderdate());
        orderDto.setOrderpriority(order.getO_orderpriority());
        orderDto.setClerk(order.getO_clerk());
        orderDto.setShippriority(order.getO_shippriority());
        orderDto.setComment(order.getO_comment());

        CustomerDto customerDto = new CustomerDto();
        Customer customer = customerHashMap.get(order.getO_custkey());
        customerDto.setCustkey(customer.getC_custkey());
        customerDto.setName(customer.getC_name());
        customerDto.setAddress(customer.getC_address());
        customerDto.setPhone(customer.getC_phone());
        customerDto.setAcctbal(customer.getC_acctbal());
        customerDto.setMktsegment(customer.getC_mktsegment());
        customerDto.setComment(customer.getC_comment());

        NationDto nationDto = new NationDto();
        Nation nation = nationHashMap.get(customer.getC_nationkey());
        nationDto.setNationkey(nation.getN_nationkey());
        nationDto.setName(nation.getN_name());
        nationDto.setComment(nation.getN_comment());

        RegionDto regionDto = new RegionDto();
        Region region = regionHashMap.get(nation.getN_regionkey());
        regionDto.setRegionkey(region.getR_regionkey());
        regionDto.setName(region.getR_name());
        regionDto.setComment(region.getR_comment());

        nationDto.setRegion(regionDto);
        customerDto.setNation(nationDto);
        orderDto.setCustomer(customerDto);

        PartSuppDto partSuppDto = new PartSuppDto();
        PartSupp partSupp = partSuppHashMap.get(new Tuple(lineItem.getL_partkey(), lineItem.getL_suppkey()));
        partSuppDto.setAvailqty(partSupp.getPs_availqty());
        partSuppDto.setSupplycost(partSupp.getPs_supplycost());
        partSuppDto.setComment(partSupp.getPs_comment());

        PartDto partDto = new PartDto();
        Part part = partHashMap.get(partSupp.getPs_partkey());
        partDto.setPartkey(part.getP_partkey());
        partDto.setName(part.getP_name());
        partDto.setMfgr(part.getP_mfgr());
        partDto.setBrand(part.getP_brand());
        partDto.setType(part.getP_type());
        partDto.setSize(part.getP_size());
        partDto.setContainer(part.getP_container());
        partDto.setRetailprice(part.getP_retailprice());
        partDto.setComment(part.getP_comment());

        SupplierDto supplierDto = new SupplierDto();
        Supplier supplier = supplierHashMap.get(partSupp.getPs_suppkey());
        supplierDto.setSuppkey(supplier.getS_suppkey());
        supplierDto.setName(supplier.getS_name());
        supplierDto.setAddress(supplier.getS_address());
        supplierDto.setPhone(supplier.getS_phone());
        supplierDto.setAcctbal(supplier.getS_acctbal());
        supplierDto.setComment(supplier.getS_comment());

        nationDto = new NationDto();
        nation = nationHashMap.get(supplier.getS_nationkey());
        nationDto.setNationkey(nation.getN_nationkey());
        nationDto.setName(nation.getN_name());
        nationDto.setComment(nation.getN_comment());

        regionDto = new RegionDto();
        region = regionHashMap.get(nation.getN_regionkey());
        regionDto.setRegionkey(region.getR_regionkey());
        regionDto.setName(region.getR_name());
        regionDto.setComment(region.getR_comment());

        nationDto.setRegion(regionDto);
        supplierDto.setNation(nationDto);
        partSuppDto.setPart(partDto);
        partSuppDto.setSupplier(supplierDto);

        lineItemDto.setOrder(orderDto);
        lineItemDto.setPartsupp(partSuppDto);

        return lineItemDto;
    }

//    private List<LineItemDto> data;
//    private void refactorSchema() {
//        data = new ArrayList<>();
//        for (LineItem lineItem : lineItemHashMap.values()) {
//            LineItemDto lineItemDto = transLineItem(lineItem);
//            data.add(lineItemDto);
//        }
//    }

    private List<LineItemDto> refactorSchema(List<LineItem> lineItems) {
        List<LineItemDto> results = new ArrayList<>(lineItems.size() + 1);
        for (LineItem lineItem : lineItems) {
            results.add(transLineItem(lineItem));
        }
        return results;
    }

    public class BatchInsert extends Thread {
        private List<DBObject> lineItems;

        public BatchInsert(List<DBObject> lineItems) {
            super();
            this.lineItems = lineItems;
        }

        @Override
        public void run() {
//            System.out.println("insert a lineItem batch");
//            AdvancedDatastore datastore = (AdvancedDatastore) morphia.createDatastore(mongoClient, "deals");
//            datastore.insert(refactorSchema(lineItems));
            mongoClient.getDB("tpch").getCollection("deals").insert(lineItems);
            lineItems.clear();
            lineItems = null;
        }
    }

    private void loadLineItemTable() {
        String tableName = "LineItem";
        CastTypes[] types = lineitemTypes;
        BufferedReader bufferedReader = null;
//        List<LineItem> lineItemList = new ArrayList<>(configCommitCount + 1);
        List<DBObject> lineItemDtoList = new ArrayList<>(configCommitCount + 1);
//        List<Thread> threadPool = new ArrayList<>();
        BatchInsert batchInsert = null;
        int recordsRead = 0;

        try {
            File file = new File(DBtool.getDBpath(getScaleFactor())
                    , tableName.toLowerCase() + ".tbl");
            bufferedReader = new BufferedReader(new InputStreamReader(new FileInputStream(file)));
            String line;
            while ((line = bufferedReader.readLine()) != null) {
                LineItem item = new LineItem();
                String[] fields = line.split("\\|");
                for (int i = 0; i < types.length; i++) {
                    String field = fields[i];
                    switch (types[i]) {
                        case LONG:
                            item.setInt(i + 1, Integer.parseInt(field));
                            break;
                        case DOUBLE:
                            item.setDouble(i + 1, Double.parseDouble(field));
                            break;
                        case STRING:
                            item.setString(i + 1, field);
                            break;
                        case DATE:
                            // yyyy-mm-dd
                            Pattern isoFmt = Pattern.compile("^\\s*(\\d{4})-(\\d{2})-(\\d{2})\\s*$");
                            Matcher isoMatcher = isoFmt.matcher(field);
                            Date date = null;
                            if (isoMatcher.find()) {
                                date = new Date(Integer.parseInt(isoMatcher.group(1)) - 1900,
                                        Integer.parseInt(isoMatcher.group(2)),
                                        Integer.parseInt(isoMatcher.group(3)));
                            } else {
                                throw new RuntimeException("Unrecognized date \""
                                        + field + "\" in TBL file: "
                                        + file.getAbsolutePath());
                            }
                            item.setDate(i + 1, date);
                            break;
                        default:
                            throw new RuntimeException("Unrecognized type for prepared statement");
                    }
                }
//                data.put((K) item.getPrimaryKey(), item);
//                lineItemList.add(item);
                lineItemDtoList.add(morphia.toDBObject(transLineItem(item)));
                recordsRead++;
                if (recordsRead % configCommitCount == 0) {
                    if (batchInsert != null) {
                        try {
                            batchInsert.join();
                        } catch (InterruptedException e) {
                            e.printStackTrace();
                        }
                    }
                    batchInsert = new BatchInsert(lineItemDtoList);
                    batchInsert.start();
//                    Thread batchInsert = new Thread(new BatchInsert(lineItemList));
//                    threadPool.add(batchInsert);
//                    batchInsert.start();
                    lineItemDtoList = new ArrayList<>(configCommitCount + 1);
                }
            }
            if (lineItemDtoList.size() > 0) {
                batchInsert = new BatchInsert(lineItemDtoList);
                batchInsert.start();
                try {
                    batchInsert.join();
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
//                Thread batchInsert = new Thread(new BatchInsert(lineItemList));
//                threadPool.add(batchInsert);
//                batchInsert.start();
            }
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (bufferedReader != null){
                try {
                    bufferedReader.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }

//        for (Thread thread : threadPool) {
//            try {
//                thread.join();
//            } catch (InterruptedException e) {
//                e.printStackTrace();
//            }
//        }
        totalRows += recordsRead;
    }

    private long totalRows = 0;
    private HashMap<Tuple, LineItem> lineItemHashMap;
    private HashMap<Long, Customer> customerHashMap;
    private HashMap<Long, Nation> nationHashMap;
    private HashMap<Long, Orders> ordersHashMap;
    private HashMap<Long ,Part> partHashMap;
    private HashMap<Tuple, PartSupp> partSuppHashMap;
    private HashMap<Long, Region> regionHashMap;
    private HashMap<Long, Supplier> supplierHashMap;

    private void loadTBL() {
        TableLoader[] tableLoaders = new TableLoader[8];
//        tableLoaders[0] = new TableLoader<LineItem, Tuple>("LineItem", lineitemTypes, this, new LineItemConstructor());
        tableLoaders[1] = new TableLoader<Customer, Long>("Customer", customerTypes, this, new CustomerConstructor(), (int) (scaleFactor * 150000));
        tableLoaders[2] = new TableLoader<Nation, Long>("Nation", nationTypes, this, new NationConstructor(), 25);
        tableLoaders[3] = new TableLoader<Orders, Long>("Orders", ordersTypes, this, new OrdersConstructor(), (int) (scaleFactor * 1500000));
        tableLoaders[4] = new TableLoader<Part, Long>("Part", partTypes, this, new PartConstructor(), (int) (scaleFactor * 200000));
        tableLoaders[5] = new TableLoader<PartSupp, Tuple>("PartSupp", partsuppTypes, this, new PartSuppConstructor(), (int) (scaleFactor * 800000));
        tableLoaders[6] = new TableLoader<Region, Long>("Region", regionTypes, this, new RegionConstructor(), 5);
        tableLoaders[7] = new TableLoader<Supplier, Long>("Supplier", supplierTypes, this, new SupplierConstructor(), (int) (scaleFactor * 10000));

        Thread loaders[] = new Thread[8];
        for (int i = 0; i < 8; i++) {
            if (tableLoaders[i] != null) {
                loaders[i] = new Thread(tableLoaders[i]);
                loaders[i].start();
            }
        }

        for (int i = 0; i < 8; i++) {
            try {
                if (loaders[i] != null)
                    loaders[i].join();
            } catch(InterruptedException e) {
                e.printStackTrace();
            }
        }

//        lineItemHashMap = tableLoaders[0].getData();
        customerHashMap = tableLoaders[1].getData();
        nationHashMap = tableLoaders[2].getData();
        ordersHashMap = tableLoaders[3].getData();
        partHashMap = tableLoaders[4].getData();
        partSuppHashMap = tableLoaders[5].getData();
        regionHashMap = tableLoaders[6].getData();
        supplierHashMap = tableLoaders[7].getData();
    }

    private class TableLoader<T extends TPCHDomain, K> implements Runnable{
        private String tableName;
        private CastTypes[] types;
        private TPCHLoaderMongoFast parent;
        private Constructor<T> constructor;
        private HashMap<K, T> data;

        TableLoader(String tableName, CastTypes[] types, TPCHLoaderMongoFast parent,
                    Constructor<T> constructor, int initCapacity) {
            this.tableName = tableName;
            this.types = types;
            this.parent = parent;
            this.constructor = constructor;
            this.data = new HashMap<>((int) (initCapacity * 1.1));
        }

        public HashMap<K, T> getData() {
            return data;
        }

        @Override
        public void run() {
            BufferedReader bufferedReader = null;
            int recordsRead = 0;

            try {
                File file = new File(DBtool.getDBpath(getScaleFactor())
                        , tableName.toLowerCase() + ".tbl");
                bufferedReader = new BufferedReader(new InputStreamReader(new FileInputStream(file)));
                String line;
                while ((line = bufferedReader.readLine()) != null) {
                    T item = constructor.newInstance();
                    String[] fields = line.split("\\|");
                    for (int i = 0; i < types.length; i++) {
                        String field = fields[i];
                        switch (types[i]) {
                            case LONG:
                                item.setInt(i + 1, Integer.parseInt(field));
                                break;
                            case DOUBLE:
                                item.setDouble(i + 1, Double.parseDouble(field));
                                break;
                            case STRING:
                                item.setString(i + 1, field);
                                break;
                            case DATE:
                                // yyyy-mm-dd
                                Pattern isoFmt = Pattern.compile("^\\s*(\\d{4})-(\\d{2})-(\\d{2})\\s*$");
                                Matcher isoMatcher = isoFmt.matcher(field);
                                Date date = null;
                                if (isoMatcher.find()) {
                                    date = new Date(Integer.parseInt(isoMatcher.group(1)) - 1900,
                                            Integer.parseInt(isoMatcher.group(2)),
                                            Integer.parseInt(isoMatcher.group(3)));
                                } else {
                                    throw new RuntimeException("Unrecognized date \""
                                            + field + "\" in TBL file: "
                                            + file.getAbsolutePath());
                                }
                                item.setDate(i + 1, date);
                                break;
                            default:
                                throw new RuntimeException("Unrecognized type for prepared statement");
                        }
                    }
                    data.put((K) item.getPrimaryKey(), item);
                    recordsRead++;
                }
            } catch (IOException e) {
                e.printStackTrace();
            } finally {
                if (bufferedReader != null){
                    try {
                        bufferedReader.close();
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                }
            }
            synchronized (parent) {
                parent.totalRows += recordsRead;
            }
        }

    }

    private abstract class Constructor<T> {
        public abstract T newInstance();
    }

    private class CustomerConstructor extends Constructor<Customer> {
        @Override
        public Customer newInstance() {
            return new Customer();
        }
    }
    private class LineItemConstructor extends Constructor<LineItem> {
        @Override
        public LineItem newInstance() {
            return new LineItem();
        }
    }
    private class NationConstructor extends Constructor<Nation> {
        @Override
        public Nation newInstance() {
            return new Nation();
        }
    }
    private class OrdersConstructor extends Constructor<Orders> {
        @Override
        public Orders newInstance() {
            return new Orders();
        }
    }
    private class PartConstructor extends Constructor<Part> {
        @Override
        public Part newInstance() {
            return new Part();
        }
    }
    private class PartSuppConstructor extends Constructor<PartSupp> {
        @Override
        public PartSupp newInstance() {
            return new PartSupp();
        }
    }
    private class RegionConstructor extends Constructor<Region> {
        @Override
        public Region newInstance() {
            return new Region();
        }
    }
    private class SupplierConstructor extends Constructor<Supplier> {
        @Override
        public Supplier newInstance() {
            return new Supplier();
        }
    }

    private static class TestLoadTBL {
        public static void main(String[] args) {
            TPCHLoaderMongoFast loader = new TPCHLoaderMongoFast(0.01);
            loader.loadTBL();
//            loader.refactorSchema();
            System.out.println("OK");
        }
    }

}
